import random
import numpy as np
from scipy import ndimage
from scipy.ndimage.interpolation import zoom
import torch
from .image_transform import make_square

def random_rot_flip(image, label):
    k = np.random.randint(0, 4)
    image = np.rot90(image, k)
    label = np.rot90(label, k)
    axis = np.random.randint(0, 2)
    image = np.flip(image, axis=axis).copy()
    label = np.flip(label, axis=axis).copy()
    return image, label


def random_rotate(image, label):
    angle = np.random.randint(-20, 20)
    image = ndimage.rotate(image, angle, order=0, reshape=False)
    label = ndimage.rotate(label, angle, order=0, reshape=False)
    return image, label

def random_clip(image, label, size):
    """
    Randomly clip a section from the given image and label.

    Parameters:
    image (np.array): The input image array. 
    label (np.array): Associated label for the image.
    size (list): Desired clip size as (new_h, new_w).

    Returns:
    Tuple: A tuple containing the clipped image and label.
    """

    assert image.shape[:2] == label.shape[:2]
    h, w = image.shape[:2]
    new_h, new_w = size[0], size[1]

    # If desired clip size is larger than the original image, adjust to fit within the image
    if h < new_h or w < new_w:
        new_h = new_w = min(h, w)

    top = np.random.randint(0, h - new_h)
    left = np.random.randint(0, w - new_w)

    image = image[top:top + new_h, left:left + new_w]
    label = label[top:top + new_h, left:left + new_w]

    return image, label


class RandomGenerator(object):
    def __init__(self, output_size, clip=False):
        self.output_size = output_size
        self.clip = clip

    def __call__(self, data):
        image, label = data
        if not self.clip:
            # zero padding to make the image square
            image, label = make_square(image), make_square(label)
        else:
            # random clip a square from the image, size=self.output_size is not ensured
            image, label = random_clip(image, label, size=self.output_size)

        if random.random() > 0.5:
            image, label = random_rot_flip(image, label)
        elif random.random() > 0.5:
            image, label = random_rotate(image, label)
        shape = image.shape
        if len(shape)==3:
            x, y, _ = shape
        else:
            x, y = shape
        # zoom image to output_size*output_size
        if x != self.output_size[0] or y != self.output_size[1]:
            if len(shape)==3:
                image = zoom(image, (self.output_size[0] / x, self.output_size[1] / y, 1), order=3)  # why not 3?
            else:
                image = zoom(image, (self.output_size[0] / x, self.output_size[1] / y), order=3)  # why not 3?
            
            label = zoom(label, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        image = torch.from_numpy(image.astype(np.float32))
        if len(shape)==3:
            image = image.permute(2,0,1) # channel becomes 1st dimension
        else:
            image = image.unsqueeze(0) # if gray scale, add channel dimension
        label = torch.from_numpy(label.astype(np.float32))
        return image, label.long()
